"""
AI自动分类服务
使用Ollama LLM分析文档内容并自动分类
"""
import logging
from typing import List, Dict, Optional, Tuple
import httpx
from sqlalchemy.orm import Session

from app.core.config import settings
from app.models.document import Document
from app.models.category import Category

logger = logging.getLogger(__name__)


class AIClassifier:
    """AI文档分类器"""
    
    def __init__(self):
        self.ollama_base_url = settings.OLLAMA_BASE_URL
        self.default_model = settings.OLLAMA_LLM_MODEL
        self.timeout = 60.0  # 60秒超时
    
    async def classify_document(
        self,
        document: Document,
        db: Session,
        model_name: Optional[str] = None,
        user_preferences: Optional[Dict] = None
    ) -> Tuple[Optional[str], float, str]:
        """
        对单个文档进行智能分类
        
        Args:
            document: 文档对象
            db: 数据库会话
            model_name: 指定使用的模型（可选）
            user_preferences: 用户偏好设置（可选）
        
        Returns:
            (category_id, confidence, reasoning) - 分类ID、置信度、推理过程
        """
        try:
            # 1. 获取所有可用分类
            categories = db.query(Category).all()
            if not categories:
                logger.warning("系统中没有可用的分类")
                return None, 0.0, "系统中没有可用的分类"
            
            # 2. 准备文档内容摘要
            content_summary = self._prepare_content_summary(document)
            
            # 3. 准备分类列表
            category_list = "\n".join([
                f"- {cat.name}: {cat.description or '无描述'}"
                for cat in categories
            ])
            
            # 4. 构建提示词
            prompt = self._build_classification_prompt(
                content_summary,
                category_list,
                document.title
            )
            
            # 5. 调用LLM进行分类
            model = model_name or user_preferences.get('preferred_llm_model') if user_preferences else self.default_model
            response_text = await self._call_ollama(prompt, model)
            
            # 6. 解析LLM响应
            category_id, confidence, reasoning = self._parse_llm_response(
                response_text,
                categories
            )
            
            logger.info(
                f"文档 {document.id} 分类结果: "
                f"category_id={category_id}, confidence={confidence:.2f}"
            )
            
            return category_id, confidence, reasoning
            
        except Exception as e:
            logger.error(f"AI分类失败: {str(e)}", exc_info=True)
            return None, 0.0, f"分类失败: {str(e)}"
    
    async def batch_classify_documents(
        self,
        documents: List[Document],
        db: Session,
        model_name: Optional[str] = None,
        user_preferences: Optional[Dict] = None
    ) -> List[Dict]:
        """
        批量对文档进行智能分类
        
        Args:
            documents: 文档列表
            db: 数据库会话
            model_name: 指定使用的模型（可选）
            user_preferences: 用户偏好设置（可选）
        
        Returns:
            分类结果列表
        """
        results = []
        
        for document in documents:
            category_id, confidence, reasoning = await self.classify_document(
                document,
                db,
                model_name,
                user_preferences
            )
            
            results.append({
                "document_id": document.id,
                "document_title": document.title,
                "category_id": category_id,
                "confidence": confidence,
                "reasoning": reasoning
            })
        
        return results
    
    def _prepare_content_summary(self, document: Document) -> str:
        """
        准备文档内容摘要（用于LLM分析）
        
        Args:
            document: 文档对象
        
        Returns:
            内容摘要字符串
        """
        # 优先使用文档摘要
        if hasattr(document, 'summary') and document.summary:
            return document.summary[:500]
        
        # 其次使用提取的文本内容
        if hasattr(document, 'extracted_text') and document.extracted_text:
            # 取前500个字符
            return document.extracted_text[:500]
        
        # 如果都没有，使用文件名和标题
        content = f"文档标题: {document.title}\n文件名: {document.file_name}"
        if hasattr(document, 'tags') and document.tags:
            try:
                import json
                tags = json.loads(document.tags) if isinstance(document.tags, str) else document.tags
                if tags:
                    content += f"\n标签: {', '.join(tags)}"
            except:
                pass
        
        return content
    
    def _build_classification_prompt(
        self,
        content_summary: str,
        category_list: str,
        document_title: str
    ) -> str:
        """
        构建分类提示词
        
        Args:
            content_summary: 文档内容摘要
            category_list: 分类列表
            document_title: 文档标题
        
        Returns:
            完整的提示词
        """
        prompt = f"""你是一个专业的文档分类助手。请根据文档的标题和内容摘要，从给定的分类列表中选择最合适的分类。

**文档标题**: {document_title}

**文档内容摘要**:
{content_summary}

**可用分类列表**:
{category_list}

**任务要求**:
1. 仔细分析文档的标题和内容
2. 从上述分类列表中选择最合适的一个分类
3. 给出你的置信度（0-100的整数）
4. 简要说明你的分类理由

**输出格式**（严格按照以下格式）:
分类: <分类名称>
置信度: <0-100的整数>
理由: <简要说明>

请开始分析:"""
        
        return prompt
    
    async def _call_ollama(self, prompt: str, model: str) -> str:
        """
        调用Ollama API
        
        Args:
            prompt: 提示词
            model: 模型名称
        
        Returns:
            LLM响应文本
        """
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(
                    f"{self.ollama_base_url}/api/generate",
                    json={
                        "model": model,
                        "prompt": prompt,
                        "stream": False,
                        "options": {
                            "temperature": 0.3,  # 较低的温度，更确定的输出
                            "top_p": 0.9,
                        }
                    }
                )
                response.raise_for_status()
                result = response.json()
                return result.get("response", "")
        
        except Exception as e:
            logger.error(f"调用Ollama失败: {str(e)}")
            raise Exception(f"调用AI模型失败: {str(e)}")
    
    def _parse_llm_response(
        self,
        response_text: str,
        categories: List[Category]
    ) -> Tuple[Optional[str], float, str]:
        """
        解析LLM响应，提取分类结果
        
        Args:
            response_text: LLM响应文本
            categories: 可用分类列表
        
        Returns:
            (category_id, confidence, reasoning)
        """
        try:
            # 初始化返回值
            category_name = None
            confidence = 0.0
            reasoning = ""
            
            # 按行解析响应
            lines = response_text.strip().split('\n')
            for line in lines:
                line = line.strip()
                
                if line.startswith('分类:') or line.startswith('分类：'):
                    category_name = line.split(':', 1)[1].strip() if ':' in line else line.split('：', 1)[1].strip()
                
                elif line.startswith('置信度:') or line.startswith('置信度：'):
                    confidence_str = line.split(':', 1)[1].strip() if ':' in line else line.split('：', 1)[1].strip()
                    # 提取数字
                    import re
                    match = re.search(r'\d+', confidence_str)
                    if match:
                        confidence = float(match.group()) / 100.0  # 转换为0-1
                
                elif line.startswith('理由:') or line.startswith('理由：'):
                    reasoning = line.split(':', 1)[1].strip() if ':' in line else line.split('：', 1)[1].strip()
            
            # 如果没有解析出理由，使用整个响应
            if not reasoning:
                reasoning = response_text[:200]
            
            # 匹配分类
            if category_name:
                # 尝试精确匹配
                for cat in categories:
                    if cat.name == category_name:
                        return cat.id, confidence, reasoning
                
                # 尝试模糊匹配（包含关系）
                for cat in categories:
                    if category_name in cat.name or cat.name in category_name:
                        return cat.id, confidence * 0.9, f"{reasoning} (模糊匹配)"
            
            # 如果没有匹配到，返回置信度最低的结果
            logger.warning(f"无法匹配分类: {category_name}")
            return None, 0.0, f"未能找到匹配的分类: {category_name}"
        
        except Exception as e:
            logger.error(f"解析LLM响应失败: {str(e)}")
            return None, 0.0, f"解析失败: {str(e)}"
    
    async def suggest_new_category(
        self,
        document: Document,
        model_name: Optional[str] = None
    ) -> Dict[str, str]:
        """
        基于文档内容建议新分类
        
        Args:
            document: 文档对象
            model_name: 指定使用的模型（可选）
        
        Returns:
            建议的分类信息 {name, description}
        """
        try:
            content_summary = self._prepare_content_summary(document)
            
            prompt = f"""基于以下文档，建议一个新的分类名称和描述：

**文档标题**: {document.title}

**文档内容摘要**:
{content_summary}

**输出格式**（严格按照以下格式）:
分类名称: <简短的分类名称>
分类描述: <详细的分类描述>

请开始分析:"""
            
            model = model_name or self.default_model
            response_text = await self._call_ollama(prompt, model)
            
            # 解析响应
            name = ""
            description = ""
            
            lines = response_text.strip().split('\n')
            for line in lines:
                line = line.strip()
                if line.startswith('分类名称:') or line.startswith('分类名称：'):
                    name = line.split(':', 1)[1].strip() if ':' in line else line.split('：', 1)[1].strip()
                elif line.startswith('分类描述:') or line.startswith('分类描述：'):
                    description = line.split(':', 1)[1].strip() if ':' in line else line.split('：', 1)[1].strip()
            
            return {
                "name": name or "未命名分类",
                "description": description or "AI建议的新分类"
            }
        
        except Exception as e:
            logger.error(f"建议新分类失败: {str(e)}")
            return {
                "name": "建议失败",
                "description": str(e)
            }

